import pretty_midi

def get_notes_with_duration(midi_file):
    # 加载 MIDI 文件
    midi_data = pretty_midi.PrettyMIDI(midi_file)
    notes_info = []

    # 遍历所有音轨
    for instrument in midi_data.instruments:
        # 只考虑具有音高的音轨
        if not instrument.is_drum:  # 过滤掉打击乐器
            for note in instrument.notes:
                note_info = {
                    'start_time': note.start,
                    'end_time': note.end,
                    'pitch': note.pitch,
                    'pitch_name': pretty_midi.note_number_to_name(note.pitch),  # 获取音符名称
                    'duration': note.end - note.start
                }
                notes_info.append(note_info)

    return notes_info

def classify_notes(notes_info):
    # 分别按起始时间和结束时间排序
    start_sorted_notes = sorted(notes_info, key=lambda x: x['start_time'])
    end_sorted_notes = sorted(notes_info, key=lambda x: x['end_time'])

    time_list = sorted(set(note['start_time'] for note in notes_info) | set(note['end_time'] for note in notes_info))

    active_notes = []  # 当前活动的音符
    last_active_notes = []  # 记录上一次的活动音符
    results = []  # 用于存储和弦信息

    # 遍历时间列表进行和弦处理
    for current_time in time_list:
        # 将早于等于当前时间的音符加入活动音符列表
        for note in start_sorted_notes:
            if note['end_time'] <= current_time:
                continue
            if note['start_time'] <= current_time:
                if note not in active_notes:
                    active_notes.append(note)
            if note['start_time'] > current_time:
                break

        # 将结束时间早于当前时间的音符移出活动音符列表
        active_notes = [note for note in active_notes if note['start_time'] <= current_time < note['end_time']]

        # 当活动音符发生变化时，输出和弦信息
        if active_notes != last_active_notes and active_notes:  # 确保不输出空和弦
            chord_pitches = {note['pitch_name'] for note in active_notes}
            chord_info = {
                'start_time': current_time,
                'end_time': min(note['end_time'] for note in active_notes),  # 和弦的结束时间为活动音符中最早的结束时间
                'chord': '+'.join(sorted(chord_pitches)),  # 确保和弦音高的顺序一致
            }
            results.append(chord_info)  # 收集和弦信息
            last_active_notes = active_notes.copy()  # 更新上次活动音符

        for note in start_sorted_notes:
            if note['end_time'] <= current_time:
                start_sorted_notes.remove(note)
            if note['start_time'] > current_time:
                break

    return results, list(active_notes) # 返回和弦信息、当前活动音符和所有单音信息

def write_to_py_file(output_py_name='results_output.py', output_data=[], write_mode='num'):
    with open(output_py_name, "w", encoding="utf-8") as f:
        f.write("res_list = [\n")

        for item in output_data:
            if write_mode == 'num':
                formatted_item = [round(float(value), 5) for value in item]
                f.write(f"  {formatted_item},\n")
            else:
                f.write(f"  {item},\n")

        f.write("]\n")
    print(f"[INFO] Results saved to {output_py_name}.")

def main(midi_file, output_file_name='results_output.py'):
    notes_info = get_notes_with_duration(midi_file)
    chords_info, active_notes = classify_notes(notes_info)  # 解包返回值

    # 准备输出数据
    output_data = []
    for chord_info in chords_info:
        output_data.append([int(chord_info['start_time'] * 1000 + 1), int(chord_info['end_time'] * 1000 - 1), chord_info['chord']])

    # 在写入文件之前按起始时间排序
    output_data.sort(key=lambda x: x[0])

    # 将结果写入文件
    write_to_py_file(output_py_name=output_file_name, output_data=output_data, write_mode='string')

if __name__ == "__main__":
    # 输入你的 MIDI 文件路径
    midi_file_path = 'jianghao.mid'
    main(midi_file_path)